# -*- coding: utf-8 -*-
from xinanjiang.Confluence import Confluence
from xinanjiang.Evapotranspiration import Evapotranspiration
from xinanjiang.Muskingum import Muskingum
from xinanjiang.Parameter import Parameter, State
from xinanjiang.Runoff import Runoff
from xinanjiang.Source import Source
from xinanjiang.Watershed import IO, Watershed

if __name__ == '__main__':
    Path = "xinanjiang"

    parameter = Parameter()  # 流域参数
    parameter.ReadFromFile(Path)  # 读取参数文件

    io = IO()  # 流域输入与输出
    io.ReadFromFile(Path)  # 从文本中导入降雨和蒸发数据，输出流量过程到文本中

    watershed = Watershed()  # 流域分单元
    watershed.ReadFromFile(Path)  # 读取流域分块信息
    watershed.calculate(io)  # 计算各单元流域各时段降雨和水面蒸发量

    state = []  # 每个单元流域的状态
    for nw in range(int(watershed.GetnW())):
        _state = State()
        _state.ReadFromFile(Path)  # 读取模型计算时段长，h
        state.append(_state)
    print("state", state)

    evapotranspiration = Evapotranspiration()  # 流域蒸散发
    evapotranspiration.SetParmameter(parameter)  # 设置参数

    runoff = Runoff()  # 流域产流
    runoff.SetParmameter(parameter)  # 设置参数

    source = Source()  # 流域分水源
    source.SetParmameter(parameter)  # 设置参数

    confluence = Confluence()  # 呈村流域汇流
    confluence.SetParmameter(parameter)  # 设置参数

    muskingum = Muskingum()  # 呈村流域河道汇流
    muskingum.SetParmameter(parameter)  # 设置参数

    # 逐时段逐单元流域计算
    nT = io.nrows  # nT表示时段数
    nW = watershed.GetnW()  # nW为单元流域个数

    # 创建动态一维数组存储流域出口断面流量过程
    io.m_Q = list(range(int(nT)))
    for nt in range(int(nT)):
        state[0].m_Q = 0.0  # m_Q置为0
        for nw in range(int(nW)):
            state[nw].SetInput(nt, nw, watershed)

            evapotranspiration.SetState(state[nw])
            evapotranspiration.calculate()
            evapotranspiration.UpdateState(state[nw])

            runoff.SetState(state[nw])
            runoff.calculate()
            runoff.UpdateState(state[nw])

            source.SetState(state[nw])
            source.calculate()
            source.UpdateState(state[nw])

            confluence.SetState(state[nw])
            confluence.calculate()
            confluence.UpdateState(state[nw])

            muskingum.SetState(state[nw])
            muskingum.calculate()
            muskingum.UpdateState(state[nw])
        io.m_Q[nt] = state[0].m_Q
    # 输出流域出口断面流量过程到文本Q.txt中
    io.WriteToFile(Path)
